package cert

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"github.com/hashicorp/golang-lru"
	"github.com/pkg/errors"
	"log"
	"math/big"
	"net"
	"os"
	"time"
)

var (
	cert = `-----BEGIN CERTIFICATE-----
MIIDWjCCAkKgAwIBAgIQBvAXEcqKzsT5r6btGFHT0zANBgkqhkiG9w0BAQsFADBU
MQswCQYDVQQGEwJDTjELMAkGA1UECBMCU0MxDjAMBgNVBAoTBVgtTWFuMRgwFgYD
VQQLEw9YLU1hbiBTU0wgUHJveHkxDjAMBgNVBAMTBVgtTWFuMB4XDTIwMDgyNDA4
MTAyOVoXDTIxMDgyNDA4MTAyOVowVDELMAkGA1UEBhMCQ04xCzAJBgNVBAgTAlND
MQ4wDAYDVQQKEwVYLU1hbjEYMBYGA1UECxMPWC1NYW4gU1NMIFByb3h5MQ4wDAYD
VQQDEwVYLU1hbjCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAN75Vh/s
A5PYm3kKiVR4IbaIZsWJgXzSUU5efBVdWp1omyiwqbz0B5Vte/OKvHmwHsZmQS3f
8LP9FNQIdQXKN0zddYCo2z56O72U1AKWr3l1+p8YE48e6y0xwHkjHJMaCv11oi+z
d79IRN37fa3+WKFCJX+1wUf1rUHBzH8tL55tpqkH1iQKyUD2ZZiluO+Kux224Rtd
yt4+GKXDa2y5Uf0GP0sJCj+mx67uwok3myodO1kRQFBEeVfJFJEgJj1iR6Bbqt2u
3G9YuWSJAlcyEBtau3aNaOpQtTqUNYqs3RVw0IoJHVN4bWqx/S/PJ1yzIRlGmKSn
3L5hfFV30X9hs7ECAwEAAaMoMCYwEwYDVR0lBAwwCgYIKwYBBQUHAwEwDwYDVR0T
AQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAJ/QOMhOc6fUq2sw3GqIxpQiU
wrJ15C3UT05dHR3O72o5XVtmSV+31aB4sERn9Ep5undMUBR0oPF9wtjNCU4VS7nx
PVn1to4if0tAP27D028eMEM+8ZfRLOKnzP84wCJLmjoIQk7e7QNDmPRnJE+JE4y8
+HUFiomKdIqS4utiPGlWWephuqCbzy+bRqS29ZNK89gjHQ9Kbbf7/CAXqGRjbimh
vEo2wAmj/Seq35S+YIqFVPmtgsUw2Ij63gBF5YXKAGx8g8t1aUuyBpNXfUrBk+7+
T7vtQe/HLQ5od8qwvVXsgBBU2SG5VY1wa2I1KT4A1eczyegpAP2SlS+rT0uW5w==
-----END CERTIFICATE-----`
)

var (
	key = `-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDe+VYf7AOT2Jt5
ColUeCG2iGbFiYF80lFOXnwVXVqdaJsosKm89AeVbXvzirx5sB7GZkEt3/Cz/RTU
CHUFyjdM3XWAqNs+eju9lNQClq95dfqfGBOPHustMcB5IxyTGgr9daIvs3e/SETd
+32t/lihQiV/tcFH9a1Bwcx/LS+ebaapB9YkCslA9mWYpbjvirsdtuEbXcrePhil
w2tsuVH9Bj9LCQo/pseu7sKJN5sqHTtZEUBQRHlXyRSRICY9YkegW6rdrtxvWLlk
iQJXMhAbWrt2jWjqULU6lDWKrN0VcNCKCR1TeG1qsf0vzydcsyEZRpikp9y+YXxV
d9F/YbOxAgMBAAECggEBAL2ySv9lf7WstrjHqgZfHe95JFYyoiNbmyidYdubTJhW
goqfogjLZBhQyP9trOOqXsRUAmfQ+atnSO0LNOP0IhAW6PzIq53/yAuGLsj1AdzH
7qyjgyBkIZiVlzZ4FLthkLAKAC2s9dea2Fwytb3tgeHPsK9yJjneXcRv6FbaToTt
Fh9GvC8y88uorP4PLt2KPf2C1Gn+7GO+5A1WHNxZ+X56ETGDA73b/CAKR6dui3dM
UP82rVUQx6M0LjvJM2M5MZJ8Mh5tC7Hy8Oa3enLI7ou/XnagnGNvbcabGIEbsw8m
hPGWcj/9FrDkoSXBWwFcvE7ltiT0M9Zm1YD6oQkKowECgYEA84W+ODdnwtEswvMG
71uY02mOtprWv7oNw51PsJim2ZB5DgWW3FoCFxVCJfkz559g7h65Q6n/MGckNtrY
RB/S5HOD/FGXwc/3xLVfw+B8YoCIx2J+GpLcJSXtECo1hg6ip1yPYSwxAFl3aonl
2H/HQo2Yi4CvT5MSZhecftkiD3kCgYEA6mYPuwZx0/SkqYDQhIBFKLzp4QRXaQPG
zbaWBuNSMOfhJM2JAIxLlqryBGBO/NARNIKBEY0Gw8Unn54dNICRLv/kt/5Y6sQd
0OG9FTQxG6+VPpjr+ZCvDVgRzMeSCv9qQDUxIs6FO5lp8nOBGO2+3ZegcAJn6uoP
wyM1GospH/kCgYBKpoMLRKDALdMk588GbPPKSx9+hrSyDVzOzFhzcDgllY4maDw4
zDa70L54WY45wLe8BxsPzPr3YYbOcgO9zmudsHOTU4fL1+7CFD+JEjXPFEvZoy25
A7AMizXCsM27jJz75Pnm9d8HlmVu+VrGJVV152YNzaLwJ8uKFE7a68DvMQKBgCpv
qd9pTkhaXUlWyemDM4IpmYokeBLKHT73FHhQxXuXL3qzP197BuBTOJe0O/inTsYb
7nQX/SaKVh8eyBjOQfqpV2Y8RSrTsLY1eD8TgOSLSFDZ7Lb1C8w1qptRck5NaZYQ
S0IK87ITl0vOv4uc11ZU8jX1KY0tLQdeI2oX0a05AoGASIbSnWXarjauxSXUDCvX
kESCwlapPcoR12GYJ1YRFsffWQkxAc8gqfUg55JOg3qm0d7WGrw4DHa4GdJUAR0a
nFHDzWn72xAy0ZNbXZT6yxMncmo+0kF1Hs8ZCKo9DwpHAeioedvMRBBPP3/r4WPR
Jcjjov/bWoVDwJvRZ0dSQQk=
-----END PRIVATE KEY-----`
)

type CertManager struct {
	rootCa  *x509.Certificate
	rootKey interface{}
	cache   *lru.Cache

	rootRaw []byte
	keyRaw  []byte
}

func CertFrom() ([]byte, []byte) {
	return []byte(cert), []byte(key)
}

func NewCertManager() (*CertManager, error) {
	cm := new(CertManager)
	cm.rootRaw, cm.keyRaw = CertFrom()

	block, _ := pem.Decode(cm.rootRaw)
	rootCa, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, err
	}
	cm.rootCa = rootCa
	block, _ = pem.Decode(cm.keyRaw)
	rootKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	cm.rootKey = rootKey
	cm.cache, _ = lru.New(1024)
	return cm, nil
}

func (cm *CertManager) GenerateTlsByHost(host string) (*tls.Config, error) {
	if h, _, _ := net.SplitHostPort(host); host != "" {
		host = h
	}

	if v, ok := cm.cache.Get(host); ok {
		return v.(*tls.Config), nil
	}
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{`X-Man`},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(1, 0, 0),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}
	if ip := net.ParseIP(host); err != nil {
		template.IPAddresses = append(template.IPAddresses, ip)
	} else {
		template.DNSNames = append(template.DNSNames, host)
	}
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, cm.rootCa, &priv.PublicKey, cm.rootKey)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create certificate")
	}
	certBlock := &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: derBytes,
	}
	serverCert := pem.EncodeToMemory(certBlock)

	keyBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(priv),
	}
	serverKey := pem.EncodeToMemory(keyBlock)

	conf, err := tls.X509KeyPair(serverCert, serverKey)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load x509 key pair")
	}

	tlsConf := &tls.Config{
		Certificates: []tls.Certificate{conf},
	}

	cm.cache.Add(host, tlsConf)

	return tlsConf, nil
}

func (cm *CertManager) RootRaw() []byte {
	return cm.rootRaw
}

func (cm CertManager) NewCertKey() {
	var priv interface{}
	var err error
	priv, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		log.Fatalf("Failed to generate private key: %v", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(365 * 24 * time.Hour)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("Failed to generate serial number: %v", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization:       []string{"X-Man"},
			Country:            []string{"CN"},
			OrganizationalUnit: []string{"X-Man SSL Proxy"},
			Province:           []string{"SC"},
			CommonName:         "X-Man",
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}
	template.IsCA = true
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, priv.(*rsa.PrivateKey).PublicKey, priv)
	if err != nil {
		log.Fatalf("Failed to create certificate: %v", err)
	}

	certOut, err := os.Create("cert.pem")
	if err != nil {
		log.Fatalf("Failed to open cert.pem for writing: %v", err)
	}
	if err := pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		log.Fatalf("Failed to write data to cert.pem: %v", err)
	}
	if err := certOut.Close(); err != nil {
		log.Fatalf("Error closing cert.pem: %v", err)
	}
	log.Print("wrote cert.pem\n")

	keyOut, err := os.OpenFile("key.pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		log.Fatalf("Failed to open key.pem for writing: %v", err)
		return
	}
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		log.Fatalf("Unable to marshal private key: %v", err)
	}
	if err := pem.Encode(keyOut, &pem.Block{Type: "PRIVATE KEY", Bytes: privBytes}); err != nil {
		log.Fatalf("Failed to write data to key.pem: %v", err)
	}
	if err := keyOut.Close(); err != nil {
		log.Fatalf("Error closing key.pem: %v", err)
	}
	log.Print("wrote key.pem\n")
}
